#include <stdio.h>
#include <cuda_runtime.h>

// 定义向量长度
#define N (16 * 1024 * 1024)

// CUDA Kernel：向量加法（每个线程处理一个元素）
__global__ void vectorAdd(const float *A, const float *B, float *C, int numElements) {
	int i = blockDim.x * blockIdx.x + threadIdx.x;
	if (i < numElements) {
		C[i] = A[i] + B[i];  // 逐元素相加
	}
}

// 主函数
int main() {
	size_t size = N * sizeof(float);

	// 主机端（CPU）数据
	float *h_A, *h_B, *h_C;
	cudaMallocHost(&h_A, size);  // 使用 pinned memory
	cudaMallocHost(&h_B, size);
	cudaMallocHost(&h_C, size);

	for (int i = 0; i < N; i++) {
		h_A[i] = rand() / (float)RAND_MAX;
		h_B[i] = rand() / (float)RAND_MAX;
	}

	// 设备端（GPU）数据指针
	float *d_A, *d_B, *d_C;
	cudaMalloc((void**)&d_A, N * sizeof(float));
	cudaMalloc((void**)&d_B, N * sizeof(float));
	cudaMalloc((void**)&d_C, N * sizeof(float));

	// 定义线程块和网格大小
	int threadsPerBlock = 256;
	int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;

	// 创建计时事件
	cudaEvent_t start, stop;
	cudaEventCreate(&start);
	cudaEventCreate(&stop);

	//开始计时
	cudaEventRecord(start);

	// 数据拷贝: Host -> Device (异步传输)
	cudaStream_t stream;
	cudaStreamCreate(&stream);
	cudaMemcpyAsync(d_A, h_A, size, cudaMemcpyHostToDevice, stream);
	cudaMemcpyAsync(d_B, h_B, size, cudaMemcpyHostToDevice, stream);

	// 启动CUDA Kernel
	vectorAdd<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(d_A, d_B, d_C, N);

	// 数据拷贝: Device -> Host (异步)
	cudaMemcpyAsync(h_C, d_C, size, cudaMemcpyDeviceToHost, stream);

	// 同步流
	cudaStreamSynchronize(stream);

	//停止计时
	cudaEventRecord(stop);
	cudaEventSynchronize(stop);

	// 计算时间
	float totalMs;
	cudaEventElapsedTime(&totalMs, start, stop);
	printf("数据量: %d 个元素\n", N);
	printf("总耗时: %.3f ms\n", totalMs);

	// 验证前10个结果
	for (int i = 0; i < 10; i++) {
		printf("h_C[%d] = %.2f (expected: %.2f)\n", i, h_C[i], h_A[i] + h_B[i]);
	}

	// 释放设备内存
	cudaFree(d_A);
	cudaFree(d_B);
	cudaFree(d_C);

	// 释放主机上的内存
	cudaFreeHost(h_A);
	cudaFreeHost(h_B);
	cudaFreeHost(h_C);

	//销毁Stream
	cudaStreamDestroy(stream);

	//销毁Event
	cudaEventDestroy(start);
	cudaEventDestroy(stop);

	return 0;
}
